package org.apache.spark.deploy

import org.apache.spark.{SecurityManager, SparkConf}
import org.apache.spark.rpc.{RpcAddress, RpcEndpointRef, RpcEnv, ThreadSafeRpcEndpoint}

class MyMaster(override val rpcEnv: RpcEnv, address: RpcAddress,
               val securityMgr: SecurityManager, val conf: SparkConf
              ) extends ThreadSafeRpcEndpoint {


  override def onConnected(remoteAddress: RpcAddress): Unit = {
    super.onConnected(remoteAddress)
    println(s"$remoteAddress is connected")
  }

  override def onDisconnected(remoteAddress: RpcAddress): Unit = {
    super.onDisconnected(remoteAddress)
    println(s"$remoteAddress is disconnected")
  }

  override def onStart(): Unit = {
    println("fake Master is starting")
  }

  override def receive: PartialFunction[Any, Unit] = {
    case TestAdd(i, j) => println(s"plus answer: ${i + j}")
    case TestSub(i, j) => println(s"subtract answer: ${i - j}")
    case RegisterWorker(workerId, host, port, workerRef, cores, memory, masterAddress) =>
      println(s"RegisterWorker: workerId=$workerId workerAddress=$host:$port cores=$cores memory=$memory")
      workerRef.send(RegisteredWorker(self, masterAddress))
    case Heartbeat(workerId, worker) => println(s"Heartbeat: 接收到${workerId}心跳")
    case _ => println("Master 接收到 Unknown Message")
  }
}

object MyMaster {
  val SYSTEM_NAME = "fakeMaster"
  val ENDPOINT_NAME = "Master"

  def main(args: Array[String]): Unit = {
    val host = "localhost"
    val port = 9999
    val conf = new SparkConf
    val securityManager = new SecurityManager(conf)
    val rpcEnv = RpcEnv.create(SYSTEM_NAME, host, port, conf, securityManager)
    val masterEndpoint = rpcEnv.setupEndpoint(ENDPOINT_NAME,
      new MyMaster(rpcEnv, rpcEnv.address, securityManager, conf))

    masterEndpoint.send(TestAdd(1, 2))
    masterEndpoint.send(TestSub(100, 1))

    rpcEnv.awaitTermination()
  }
}

case class TestAdd(i: Int, j: Int)
case class TestSub(i: Int, j: Int)

case class RegisterWorker(workerId: String,
                          host: String,
                          port: Int,
                          worker: RpcEndpointRef,
                          cores: Int,
                          memory: Int,
                          masterAddress: RpcAddress)

case class RegisteredWorker(master: RpcEndpointRef,
                             masterAddress: RpcAddress)

case class Heartbeat(workerId: String, worker: RpcEndpointRef)